#include "refine/smooth_despill_known_back.h"
#include "refine/window_pass.h"
#include "refine/copy.h"
#include <limits>
#include <iostream>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>
#include <glib/g2/smatrix.h>
#include <glib/g3/smatrix.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_despill_known_back {

typedef unsigned uint;
typedef unsigned long long ulong;
typedef g3::basic_vector<uint>		vectoru;
typedef g3::basic_smatrix<uint>		smatrixu;
typedef g3::basic_vector<ulong>		vectorul;
typedef g3::basic_smatrix<ulong>	smatrixul;
typedef g3::basic_vector<float>		vectorf;
typedef g3::basic_smatrix<float>	smatrixf;
typedef g2::basic_vector<float>		vector2f;
typedef g2::basic_smatrix<float>	smatrix2f;

typedef unsigned char uchar;
const uint uchar_max = std::numeric_limits<uchar>::max();
const uint uchar_avg = uchar_max / 2;
const float k = uchar_max;
const float kk = uchar_max * uchar_max;

template <typename Color>
inline vectoru v(Color const & c)
	{ return vectoru(c.b, c.g, c.r); }

inline uchar limit(float v) {
	return (v < 0) ? 0 : (v >= uchar_max) ? uchar_max : uchar(v + 0.5f);
}

template <typename Color>
inline void discrete(vectorf const & v, Color & c) {
	c.b = limit(v.x);
	c.g = limit(v.y);
	c.r = limit(v.z);
}

template <typename Raster>
class builder
{
public:
	typedef typename Raster::elem_type elem_type;
public:
	builder(
		Raster const & im_, 
		Raster const & back_,
		Raster & next_im_, 
		img::bitmap & modifiable_, 
		float max_condition_number_, float blending_prob_)
		: im(im_), back(back_), next_im(next_im_), modifiable(modifiable_),
		  max_condition_number(max_condition_number_), blending_prob(blending_prob_)
	{
		clear(); 
	}
	void clear() {
		remove_all();
		checks = achanges = 0;
	}
	size_extent dim() const
		{ return im.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) 
		{ add_sample(v(im[p]), im[p].a); }
	void remove(size_point const & p) 
		{ remove_sample(v(im[p]), im[p].a); }
	void remove_all() {
		n = A = AA = 0;
		AAA = AAAA = 0;
		AC = vectoru::nill;
		AAC = AAAC = vectorul::nill;
		AACC.zero();
	}
	void process(size_point const & p, size_rect const &) {
		++checks;
		elem_type & e = next_im[p];
		vectoru c = v(im[p]);
		vectoru c0 = v(back[p]);
		if (c == c0) {
			if (im[p].a != 0)
				e.a = 0;
			else
				modifiable[p] = false;
		}
		else if (!classify(c, c0, e) || im[p].a == e.a)
			modifiable[p] = false;
		else
			++achanges;
	}

private:
	bool classify(vectorf const & c, vectorf const & c0, elem_type & e) const 
	{
		assert (c != c0);
		if (AA * max_condition_number < n * kk)
			return false;

		// d is vector from c0 towards mean foreground color (not towards c)
		vectorf d = k * (vectorf(AC) - float(A) * c0) / float(AA);

		// Sigma * AAAA
		smatrixf Sigma = 
			  kk * smatrixf(AACC) 
			+ kk * AA * tsquare(c0) 
			+ float(AAAA) * tsquare(d)
			- kk * sproduct(c0, vectorf(AAC))
			+ k * float(AAA) * sproduct(c0, d)
			- k * sproduct(vectorf(AAAC), d);

		float tr = Sigma.tr();
		// really trace can't be negative, since the matrix is nonnegative definite,
		// although it can happen due to round-off errors for null matrix;
		// if the matrix is null then alpha classification is perfect already
		if (tr <= 0) 
			return false;
		// normalize matrix to have its trace equal to 1
		// (otherwise det() is too large for float)
		Sigma *= 1 / tr;
		Sigma += 0.001f;

		smatrixf A = Sigma.inverse();
		vectorf d0 = c - c0;
		vectorf d1 = d0 - d;

		//classify as in smooth_despill_b
		vectorf Ad0 = A(d0);
		vectorf Ad1 = A(d1);

		smatrix2f S(
			scalar_product(d0, Ad0),
			scalar_product(d1, Ad0),
			scalar_product(d1, Ad1));
		vector2f beta = S.eigenvector(1);
		vectorf n = beta.x * d0 + beta.y * d1;

		float nAn = scalar_product(n, A(n));
		float den = 1 / nAn;
		float nAd0 = scalar_product(Ad0,n);
		float t0 = -den * nAd0;
		float nAd1 = scalar_product(Ad1,n);
		float t1 = -den * nAd1;

		if (t0 * t1 <= 0) {
			float c_in_0 = S.xx;
			float c_in_1 = S.yy;
			float c0_in_0 = S.xx + 2 * t0 * nAd0 + t0 * t0 * nAn;
			float c1_in_1 = S.yy + 2 * t1 * nAd1 + t1 * t1 * nAn;
			float sum = c0_in_0 + c1_in_1;
			if (sum > blending_prob * c_in_0 && c_in_0 < c_in_1)
				e.a = 0;
			else if (sum > blending_prob * c_in_1 && c_in_1 < c_in_0)
				e.a = uchar_max;
			else
				e.a = uchar_max * fabs(t0 / (t1 - t0)) + 0.5f;
			discrete(c + t1 * n, e);
		}
		else if (t0 == t1) {
			e.a = uchar_avg;
			discrete(c, e);
		}
		else if (t0*t0 > t1*t1) {
			e.a = uchar_max;
			discrete(c, e);
		}
		else {
			e.a = 0;
			discrete(c + t1 * n, e);
		}
		return true;
	}
private:
	void add_sample(vectoru const & c, uint a) {
		++n;
		A += a;
		uint aa = a * a;
		AA += aa;
		AAA += aa * a;
		AAAA += aa * aa;
		AC += a * c;
		AAC += aa * c;
		AAAC += aa * a * c;
		AACC += smatrixul(aa * tsquare(c));
	}
	void remove_sample(vectoru const & c, uint a) {
		--n;
		A -= a;
		uint aa = a * a;
		AA -= aa;
		AAA -= aa * a;
		AAAA -= aa * aa;
		AC -= a * c;
		AAC -= aa * c;
		AAAC -= aa * a * c;
		AACC -= smatrixul(aa * tsquare(c));
	}
private:
	Raster const & im;
	Raster const & back;
	Raster & next_im;
	img::bitmap & modifiable;
	float max_condition_number;
	float blending_prob;
private:
	uint n, A, AA;
	ulong AAA, AAAA;
	vectoru AC;
	vectorul AAC, AAAC;
	smatrixul AACC;
	uint checks, achanges;
};

template <typename Raster>
static void perform_core(
	Raster & im, 
	Raster const & back, 
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Smooth & Despill using known background" << endl;

	size_extent dim = im.dim();
	assert (dim == modifiable.dim());

	Raster despilled = im;

	if (p.passes_num == 0) {
		// no smoothing, despill only
		builder<Raster> sm(im, back, despilled, modifiable, p.max_condition_number, p.blending_prob);
		window_pass(sm, p.whs);
		copy_color(im, despilled);
	}
	else {
		for (uint pass = 0; ; ++pass) {
			if (!p.silent)
				clog << "pass " << pass+1 << '/' << p.passes_num << flush;

			builder<Raster> sm(im, back, despilled, modifiable, p.max_condition_number, p.blending_prob);
			window_pass(sm, p.whs);

			if (!p.silent)
				clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;

			if (pass+1 < p.passes_num && sm.change_count() > 0)
				copy_alpha(im, despilled);
			else {
				img::copy(im, despilled);
				break;
			}

			img::dilate8(modifiable, p.whs);
		}
	}
}

template <typename Raster>
static void perform_from_grey_mask(
	Raster & im, 
	Raster const & back, 
	const params & p)
{
	size_extent dim = im.dim();

	img::bitmap not_min(dim), not_max(dim)/*, strict_fore(dim), strict_back(dim)*/;
	for (typename Raster::const_iterator i = im.begin(); i != im.end(); ++i) {
		not_min[i] = (*i).a != 0;
		not_max[i] = (*i).a != uchar_max;
	}

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(im, back, modifiable, p);
}

// exported functions

void perform(
	byte_image & im, 
	byte_image const & back, 
	const params & p)
{
	perform_from_grey_mask(im, back, p);
}

void perform(
	argb_byte_image & im, 
	argb_byte_image const & back, 
	const params & p)
{
	perform_from_grey_mask(im, back, p);
}


} //namespace smooth_despill_known_back

} //namespace refine
